import pandas as pd
import numpy as np
import yaml
import os
from typing import Dict, Optional, Tuple
import logging

class DataLoader:
    """
    数据加载器 - 负责加载和初步验证原始数据
    """
    
    def __init__(self, config_path: str = "config/data_config.yaml"):
        self.logger = logging.getLogger(__name__)
        self.config = self._load_config(config_path)
        
    def _load_config(self, config_path: str) -> dict:
        """加载配置文件"""
        try:
            with open(config_path, "r", encoding="utf-8") as f:
                return yaml.safe_load(f)
        except FileNotFoundError:
            self.logger.warning(f"配置文件 {config_path} 未找到，使用默认配置")
            return self._get_default_config()
    
    def _get_default_config(self) -> dict:
        """获取默认配置"""
        return {
            'data_paths': {
                'raw': {
                    'movies': 'data/raw/movies.csv',
                    'ratings': 'data/raw/ratings.csv',
                    'tags': 'data/raw/tags.csv',
                    'links': 'data/raw/links.csv'
                },
                'processed': {
                    'output_dir': 'data/processed/'
                }
            },
            'processing': {
                'min_ratings_per_user': 5,
                'min_ratings_per_movie': 10,
                'test_size': 0.2,
                'random_state': 42
            }
        }
    
    def load_raw_data(self) -> Dict[str, pd.DataFrame]:
        """
        加载所有原始数据文件
        
        Returns:
            Dict[str, pd.DataFrame]: 数据字典
        """
        data = {}
        raw_paths = self.config['data_paths']['raw']
        
        for name, path in raw_paths.items():
            try:
                if os.path.exists(path):
                    data[name] = pd.read_csv(path)
                    self.logger.info(f"✅ 成功加载 {name}: {data[name].shape}")
                else:
                    self.logger.warning(f"⚠️  文件未找到: {path}")
                    data[name] = None
            except Exception as e:
                self.logger.error(f"❌ 加载 {name} 时出错: {e}")
                data[name] = None
        
        return data
    
    def validate_data(self, data: Dict[str, pd.DataFrame]) -> Dict[str, bool]:
        """
        验证数据完整性
        
        Args:
            data: 加载的数据字典
            
        Returns:
            验证结果字典
        """
        validation_results = {}
        
        # 检查必要文件
        required_files = ['movies', 'ratings']
        for file in required_files:
            validation_results[file] = data.get(file) is not None and not data[file].empty
        
        # 检查数据一致性
        if data['ratings'] is not None and data['movies'] is not None:
            movies_in_ratings = set(data['ratings']['movieId'].unique())
            movies_in_movies = set(data['movies']['movieId'].unique())
            missing_movies = movies_in_ratings - movies_in_movies
            validation_results['movie_consistency'] = len(missing_movies) == 0
            if missing_movies:
                self.logger.warning(f"在ratings中但不在movies中的电影: {len(missing_movies)}")
        
        return validation_results
    
    def get_data_summary(self, data: Dict[str, pd.DataFrame]) -> Dict[str, dict]:
        """
        获取数据摘要信息
        
        Args:
            data: 数据字典
            
        Returns:
            各数据集的摘要信息
        """
        summary = {}
        for name, df in data.items():
            if df is not None:
                summary[name] = {
                    'shape': df.shape,
                    'columns': list(df.columns),
                    'missing_values': df.isnull().sum().to_dict(),
                    'dtypes': df.dtypes.astype(str).to_dict()
                }
        return summary